from typing import Iterable, TypeVar

from zkl_aiutils_datasets.basics import Dataset, SkipExhausted, skip
from zkl_aiutils_datasets.wrapping import EmptyDatasetIterator, wrap_dataset_iterator
from .chained import ChainedDatasets, ChainedIterator, ChainedIteratorState, try_pause

AnySample = TypeVar('AnySample')


class EqualSizeChainedIterator(ChainedIterator[AnySample]):
    def __init__(self,
        children: Iterable[Iterable[AnySample]],
        child_samples_n: int,
        state: ChainedIteratorState | None = None,
    ):
        super().__init__(children, state)
        self.child_samples_n = child_samples_n

    def __skip__(self, samples_n: int):
        if samples_n <= 0:
            return

        # skip samples from current child
        try:
            skip(self._samples_iterator, samples_n)
            return
        except SkipExhausted as err:
            skipped_samples_n = err.skipped_samples_n

        # skip samples from by children
        skip_children_n = (samples_n - skipped_samples_n) // self.child_samples_n
        if skip_children_n > 0:
            try:
                skip(self._children_iterator, skip_children_n)
                skipped_children_n = skip_children_n
            except SkipExhausted as err:
                skipped_children_n = err.skipped_samples_n
            self._children_iterator_state_or_error = try_pause(self._children_iterator)

            if skipped_children_n < skip_children_n:
                self._samples_iterator = EmptyDatasetIterator()
                raise SkipExhausted(skipped_samples_n)

        skipped_samples_n += skip_children_n * self.child_samples_n
        if skipped_samples_n == samples_n:
            self._samples_iterator = EmptyDatasetIterator()
            return

        # skip samples from new child
        try:
            child = next(self._children_iterator)
        except StopIteration:
            self._samples_iterator = EmptyDatasetIterator()
            raise SkipExhausted(skipped_samples_n)
        self._samples_iterator = wrap_dataset_iterator(child)

        try:
            skip(self._samples_iterator, samples_n - skipped_samples_n)
        except SkipExhausted as err:
            raise ValueError(
                f"Expected any child has {self.child_samples_n} samples, "
                f"got only {err.skipped_samples_n} samples!")


class EqualSizeChainedDatasets(ChainedDatasets[AnySample]):
    def __init__(self, children: Iterable[Iterable[AnySample]], child_samples_n: int):
        super().__init__(children)
        self._child_samples_n = child_samples_n

    @property
    def child_samples_n(self) -> int:
        return self._child_samples_n

    def __len__(self) -> int:
        # noinspection PyTypeChecker
        return len(self._children) * self._child_samples_n

    def __getitem__(self, index: int, /) -> AnySample:
        child_i = index // self.child_samples_n
        sample_i = index % self.child_samples_n
        # noinspection PyUnresolvedReferences
        return self._children[child_i][sample_i]

    def __iter__(self) -> EqualSizeChainedIterator[AnySample]:
        return EqualSizeChainedIterator(self._children, self.child_samples_n)

    def __resume__(self, state: ChainedIteratorState | None = None, /) -> EqualSizeChainedIterator[AnySample]:
        return EqualSizeChainedIterator(self._children, self.child_samples_n, state)
